'''
Created on 24 juin 2011

@author: EliX
'''
from multiprocessing import (Event, Lock)
from sage.combinat.backtrack import SearchForest
from sage.misc.abstract_method import abstract_method
import threading
import multiprocessing
import time
import copy
import Queue

CONSTSYNC = 10000

class SearchForest2(SearchForest):
    
    def __init__(self, roots = None, 
                 children = None, 
                 post_process = None,
                 statistic = None, 
                 reduce_function = None, 
                 reduce_init = None, 
                 algorithm = 'depth', 
                 category=None) :
        r'''
        TESTS::
            # les nombres entre 1 et 2**6
            sage: from sage.combinat.MinSync3 import SearchForest2
            sage: S = SearchForest2( [1],            # root
            ...  lambda l: [(l<<1)|0, (l<<1)|1] if l < 1<<5 else [], # children
            ...  lambda x: x,                        # post-process
            ...  lambda x: [x],                      # statistic
            ...  lambda x,y: x+y,                    # reduce_function
            ...  [] )                                # reduce_init
        '''
        SearchForest.__init__(self, roots, children, post_process, algorithm, category)
        if post_process is None :
            self.post_process = lambda x: x
        if statistic is not None :
            self.statistic = statistic
        if reduce_function is not None :
            self.reduce_function = reduce_function
        self._res = None
        self._init = copy.copy(reduce_init)

    @abstract_method
    def statistic(self, o):
        r''' TODO '''
    
    @abstract_method
    def reduce_function(self):
        r''' TODO '''
        
    def _initialisation(self):
        self._res = copy.copy(self._init)
        n = 4 #max (multiprocessing.cpu_count(), 1)
        self._queue = Queue.Queue()
        self._event = Event()
        self._varFIN = 0
        self._varTMP = 0
        self._lock = Lock()
        self._thds = [Travelator(self, 2**i) for i in range(n)]
        self._maskFIN = 0
        for i in range(n) :
            self._maskFIN = self._maskFIN | 2**i
    
    def _distrib(self):
        li = [ self.roots() ]
        while len( li ) < len( self._thds ) :
            nodes = li.pop(0)
            for n in nodes :
                update(n, self, self)
                li.append( iter(self.children(n)) )
        i = 0
        while len(li) > 0 :
            self._thds[i % len(self._thds)]._liste.append( li.pop() )
            i = i + 1
            
        
    def _start(self):
        for t in self._thds :
            t.start()
 
    def _assemblage(self):
        for t in self._thds : t.join()
        for t in self._thds :
            self._res = self.reduce_function(self._res, t._res)
    
    def run(self):
        r'''
        TESTS::
            # SGO des mots binaires de taille inferieure a 20
            sage: from sage.combinat.MinSync3 import SearchForest2
            sage: S = SearchForest2( [[]],                           # root
            ...     lambda l: [l+[0], l+[1]] if len(l) < 20 else [], # children
            ...     lambda x: x,                                     # post-process
            ...     lambda x: var('x')**len(x),                      # statistic
            ...     lambda x,y: x+y,                                 # reduce_function
            ...     0 )                                              # reduce_init
            sage: S.run()
            1048576*x^20 + 524288*x^19 + 262144*x^18 + 131072*x^17 + 65536*x^16 + 32768*x^15 + 16384*x^14 + 8192*x^13 + 4096*x^12 + 2048*x^11 + 1024*x^10 + 512*x^9 + 256*x^8 + 128*x^7 + 64*x^6 + 32*x^5 + 16*x^4 + 8*x^3 + 4*x^2 + 2*x + 1

            # SGO des permutations de taille inferieure a 10
            sage: R = SearchForest2( [[]],                           # root
            ...     lambda l: [l[:i] + [len(l)] + l[i:] for i in range(len(l)+1)] if len(l) < 8 else [], # children
            ...     lambda x: x,                                     # post-process
            ...     lambda x: var('x')**len(x),                      # statistic
            ...     lambda x,y: x+y,                                 # reduce_function
            ...     0 )                                              # reduce_init
            sage: R.run()
            40320*x^8 + 5040*x^7 + 720*x^6 + 120*x^5 + 24*x^4 + 6*x^3 + 2*x^2 + x + 1

            # SGO des sommes strictements decroissantes bornees par 1 et 20
            sage: R = SearchForest2( [([i],i, i) for i in range(1,20)],   # root
            ...    lambda (list, sum, last): [(list + [i], sum + i, i) for i in range(1,last)], # children
            ...    lambda x: x,                                   # post-process
            ...    lambda (li, sum, _): var('y')**sum,            # statistic
            ...    lambda x,y: x + y,                             # reduce_function
            ...    0 )                                            # reduce_init
            sage: R.run()
            y^190 + y^189 + y^188 + 2*y^187 + 2*y^186 + 3*y^185 + 4*y^184 + 5*y^183 + 6*y^182 + 8*y^181 + 10*y^180 + 12*y^179 + 15*y^178 + 18*y^177 + 22*y^176 + 27*y^175 + 32*y^174 + 38*y^173 + 46*y^172 + 54*y^171 + 63*y^170 + 74*y^169 + 86*y^168 + 99*y^167 + 115*y^166 + 132*y^165 + 151*y^164 + 173*y^163 + 197*y^162 + 223*y^161 + 253*y^160 + 285*y^159 + 320*y^158 + 360*y^157 + 402*y^156 + 448*y^155 + 499*y^154 + 553*y^153 + 611*y^152 + 675*y^151 + 743*y^150 + 815*y^149 + 894*y^148 + 977*y^147 + 1065*y^146 + 1161*y^145 + 1260*y^144 + 1365*y^143 + 1477*y^142 + 1594*y^141 + 1716*y^140 + 1846*y^139 + 1980*y^138 + 2119*y^137 + 2266*y^136 + 2417*y^135 + 2572*y^134 + 2735*y^133 + 2901*y^132 + 3071*y^131 + 3248*y^130 + 3427*y^129 + 3609*y^128 + 3797*y^127 + 3986*y^126 + 4176*y^125 + 4371*y^124 + 4565*y^123 + 4760*y^122 + 4957*y^121 + 5153*y^120 + 5346*y^119 + 5541*y^118 + 5732*y^117 + 5919*y^116 + 6106*y^115 + 6287*y^114 + 6462*y^113 + 6635*y^112 + 6800*y^111 + 6958*y^110 + 7111*y^109 + 7255*y^108 + 7389*y^107 + 7518*y^106 + 7636*y^105 + 7742*y^104 + 7842*y^103 + 7929*y^102 + 8004*y^101 + 8071*y^100 + 8125*y^99 + 8165*y^98 + 8197*y^97 + 8215*y^96 + 8220*y^95 + 8215*y^94 + 8197*y^93 + 8165*y^92 + 8125*y^91 + 8071*y^90 + 8004*y^89 + 7929*y^88 + 7842*y^87 + 7742*y^86 + 7636*y^85 + 7518*y^84 + 7389*y^83 + 7255*y^82 + 7111*y^81 + 6958*y^80 + 6800*y^79 + 6635*y^78 + 6462*y^77 + 6287*y^76 + 6106*y^75 + 5919*y^74 + 5732*y^73 + 5541*y^72 + 5346*y^71 + 5153*y^70 + 4957*y^69 + 4760*y^68 + 4565*y^67 + 4371*y^66 + 4176*y^65 + 3986*y^64 + 3797*y^63 + 3609*y^62 + 3427*y^61 + 3248*y^60 + 3071*y^59 + 2901*y^58 + 2735*y^57 + 2572*y^56 + 2417*y^55 + 2266*y^54 + 2119*y^53 + 1980*y^52 + 1846*y^51 + 1716*y^50 + 1594*y^49 + 1477*y^48 + 1365*y^47 + 1260*y^46 + 1161*y^45 + 1065*y^44 + 977*y^43 + 894*y^42 + 815*y^41 + 743*y^40 + 675*y^39 + 611*y^38 + 553*y^37 + 499*y^36 + 448*y^35 + 402*y^34 + 360*y^33 + 320*y^32 + 285*y^31 + 253*y^30 + 223*y^29 + 197*y^28 + 173*y^27 + 151*y^26 + 132*y^25 + 115*y^24 + 99*y^23 + 86*y^22 + 74*y^21 + 63*y^20 + 54*y^19 + 46*y^18 + 38*y^17 + 32*y^16 + 27*y^15 + 22*y^14 + 18*y^13 + 15*y^12 + 12*y^11 + 10*y^10 + 8*y^9 + 6*y^8 + 5*y^7 + 4*y^6 + 3*y^5 + 2*y^4 + 2*y^3 + y^2 + y
        '''
        # initialisation
        self._initialisation()
        # distribution
        self._distrib()
        # demarrage du parcours
        self._start()
        # attente de la fin du parcours
        self._event.wait()
        # assemblage des résultats
        self._assemblage()
        # retour de la fonction
        return self._res
    
class Travelator(threading.Thread):
    def __init__(self, obj, pos):
        threading.Thread.__init__(self)
        self.pos = pos
        self._res = copy.copy(obj._init)
        self._liste = []
        self._obj = obj
        self._run = True

    def run(self):
        i = 0
        obj = self._obj
        while self._run :
            node = None
            if len(self._liste) > 0 :
                # cas classique
                node = self._liste.pop()
            else :
                # cas liste vide
                try : node = obj._queue.get(False)
                except Queue.Empty :
                    # si la queue est vide
                    # declaration en dormeur
                    with obj._lock :
                        obj._varTMP = obj._varTMP | self.pos
                    # tant qu'il n'y a rien dans la queue
                    while node is None and (obj._varTMP & self.pos) != 0 :
                        try : 
                            node = obj._queue.get(False)
                            with obj._lock :
                                obj._varTMP = obj._varTMP ^ self.pos
                                obj._varFIN = 0
                        except Queue.Empty :
                            if obj._varTMP == obj._maskFIN :
                                with obj._lock :
                                    obj._varFIN = obj._varFIN | self.pos
                                while (obj._varTMP == obj._maskFIN
                                       and not obj._event.is_set()) :
                                    if (obj._varFIN == obj._maskFIN 
                                            and obj._varTMP == obj._maskFIN
                                            and self.pos == 1) :
                                        obj._event.set()
                                if obj._event.is_set() :
                                    self._run = False
                                    break
                                else : obj._varFIN = 0
            
            if node is None :
                continue
     
            for n in node :
                update(n, obj, self)
                it = iter( obj.children(n) )
                i = i + 1
                self._liste.append( it ) 
                if i == CONSTSYNC :
                    i = 0
                    if obj._varTMP != 0 :
                        tmp = self._liste.pop(0)
                        obj._queue.put( tmp )

def update(n, obj, th) :
    ''' S'occupe de toute la partie maj des valeurs ''reduce_function'' 
    '''
    node = obj.post_process(n)      
    if node is None :
        return

    # maj des valeurs
    th._res = obj.reduce_function(th._res, obj.statistic(node))
